import copy
import numpy as np

from GA import GA
import copy


class tsp_ga(GA):
    def __init__(self, N=4, precison=5, mutation_rate=0.1) -> None:
        super().__init__(N=N, precison=precison, mutation_rate=mutation_rate)
        self.n=1000
        self.best_value_record=[]
        self.best_x=None
        self.best_value=1000


    def get_value(self, x):
        """
        求值函数
        """
        value=0
        for i in range(self.n):
            value+=self.data[int(x[i]),int(x[i+1])]

        return value

    
    def nn(self,start_point:str):
        # data=copy.deepcopy(self.data)
        """
        最近邻居环游
        """
        res=[int(start_point)]
        point=start_point
        for _ in range(self.n-1):
            best_value=1000
            best_idx=0
            for j in range(self.n):
                if j not in res:
                    if self.data[point,j]<best_value:
                        best_value=self.data[point,j]
                        best_idx=j
                    else:
                        continue

            res.append(best_idx)
            point=best_idx

        res.append(start_point)

        s=""
        for item in res:
            s+=str(item)
        return s

    def _initiate_population(self):
        """
        初始种群
        """
        res={}
        for i in range(self.N):
            one=self.nn(i)
            res[one]=self.get_value(one)

        return res

    def _crossover_this(self, gen_1: str, gen_2):
        """
        顺序交叉方法
        """
        gen_1=gen_1[:-1]
        gen_2=gen_2[:-1]
        child_1=list('0'*self.n)
        child_2=list('0'*self.n)
        choose_range=list(range(self.n))
        p1=np.random.choice(choose_range)

        choose_range.remove(p1)
        p2=np.random.choice(choose_range)

        c1,c2=min(p1,p2),max(p1,p2)
        
        r1,r2=gen_1[c2:],gen_2[c2:]
        l1=r1+gen_1.replace(r1,"")
        l2=r2+gen_2.replace(r2,"")
        
        # l1=copy.deepcopy(gen_1)
        # l2=copy.deepcopy(gen_2)
        for i in range(c1,c2):
            l1=l1.replace(gen_2[i],"")
            l2=l2.replace(gen_1[i],"")
            child_1[i]=gen_2[i]
            child_2[i]=gen_1[i]

        point=c2
        for i in range(len(l1)):
            child_1[point]=l1[i]
            child_2[point]=l2[i]
            point+=1
            if point>=self.n:
                point=0

        child_1.append(child_1[0])
        child_2.append(child_2[0])

        return ["".join(child_1),"".join(child_2)]


    def _get_random_two_samples(self):
        """
        由于经常随机生成两个样本，所以单独写一个随机生成函数
        返回值为样本下标
        """
        init=list(range(2,self.n))

        p1=np.random.choice(init)
        init.remove(p1)

        p2=np.random.choice(init)

        return [p1,p2]

    def _mutation_this(self,samples):
        """
        随机变异变异
        """
        res=[]
        for item in samples:
            p0=np.random.rand()
            c=copy.deepcopy(item)
            if p0<self.mutation_rate:
                idx=self._get_random_two_samples()
                init=list(c)
                init[idx[0]],init[idx[1]]=init[idx[1]],init[idx[0]]

                c="".join(init)

            res.append(c)

        return res


    def _valid_samples(self,samples):
        """
        验证是否带有缺失边
        """
        for item in samples:
            value=self.get_value(item)

            if value>1000:
                return False

        
        return True

    def _update_population_this(self,samples):
        """
        更新种群
        """
        population=copy.deepcopy(self.population)

        for it in samples:
            population[it]=self.get_value(it)

        sored_population=sorted(population.items(),key=lambda x:x[1])[:self.N]


        res=dict(sored_population)

        return res
            

    def _record_best_trace(self):
        """
        迭代记录
        """
        res_value=1000
        res_x=None
        for key in self.population:
            value=self.population[key]
            if value<res_value:
                res_value=value

                res_x=key
            if value<self.best_value:
                self.best_value=value
                self.best_x=key

        
        self.best_value_record.append([res_x,res_value])



        


    def __call__(self, data, epoch=100):
        """
        逻辑执行方法
        """
        self.data=data
        self.n=data.shape[0]
        ## 初始化种群
        self.population=self._initiate_population()
        self._record_best_trace()

        for _ in range(epoch):
            ## 生成父本
            samples=self.get_samples(self.population)
            while True:
                ## 生成子代
                children_row=self._crossover_this(samples[0],samples[1])
                ## 子代变异
                children_mutation=self._mutation_this(children_row)
                ## 判断子代是否符合条件
                flag=self._valid_samples(children_mutation)

                if flag:
                    break
            ## 更新种群
            self.population=self._update_population_this(children_mutation)
            self._record_best_trace()

            # print(self.population)

            

def result(x:str):

    dict_name={"0":"W","1":"Y","2":"B","3":"R"}
    s=dict_name[x[0]]
    for it in x[1:]:
        s+="-"
        s+=dict_name[it]
        
    return s



if __name__=="__main__":
    data=np.array([[1000,10,17,15],[20,1000,19,18],[50,44,1000,22],[45,40,20,1000]])

    data1=np.array([[1000,10,17,15,0],[20,1000,19,18,0],[50,44,1000,22,0],[45,40,20,1000,0],[0,0,0,0,1000]])


    ## 闭环
    model=tsp_ga()

    model(data=data)

    print("*"*40)
    print("the best value is {} and the best path is {}".format(model.best_value,result(model.best_x)))


    ## 开环游
    model1=tsp_ga()
    model1(data=data1)

    x=model1.best_x.replace(str(model1.n-1),"")[1:]
    print("*"*40)
    print("the best value is {} and the best path is {}".format(model1.best_value,result(x)))




        
            
